#include "ted.hh"
#include "dpm.hh"

template<typename T>
T get_param(const char* arg, std::string emsg)
{
  T val;
  try
  {
    val = boost::lexical_cast<T>(arg);
  } catch (boost::bad_lexical_cast&)
  {
    err_msg(emsg);
    exit(1);
  }
  return val;
}

template<typename T>
void vec_out(std::ofstream& ofs, const std::vector<T>& vec)
{
  if (vec.size() == 0)
    return;
  ofs << vec.at(0);
  for (size_t j = 1; j < vec.size(); ++j)
  {
    ofs << "\t" << vec.at(j);
  }
}

struct sqrt_t
{
  double operator()(const double x) const
  {
    return std::sqrt(x);
  }
};

template<typename T, typename F>
void vec_out(std::ofstream& ofs, const std::vector<T>& vec, F f)
{
  if (vec.size() == 0)
    return;
  ofs << f(vec.at(0));
  for (size_t j = 1; j < vec.size(); ++j)
  {
    ofs << "\t" << f(vec.at(j));
  }
}

////////////////////////////////////////////////////////////////
struct prog_args_t
{

  explicit prog_args_t() :
      x_file(""), y_file(""), K(2), output("output"), // default args
      bmin(-2.), bmax(2.),                            // beta in [bmin, bmax]
      repeat(1), tol(1e-3), max_iter(100), sample(100), // options
      minibatch(10)
  {
  }

  std::string x_file;
  std::string y_file;
  size_t K;
  std::string output;
  double bmin;
  double bmax;
  size_t repeat;
  double tol;
  size_t max_iter;
  size_t sample;
  size_t minibatch;
};

void print_help(const char* prog, const prog_args_t& args)
{
  std::cerr << prog << " -x X.txt -y Y.txt -k K -o out" << std::endl;
  std::cerr << "         -bmin bmin -bmax bmax" << std::endl;
  std::cerr << "         -i iter -r repeat -t tol -s sample" << std::endl;
#ifdef ONLINE
  std::cerr << "         -minibatch size" << std::endl;
#endif
  std::cerr << std::endl;
  std::cerr << " X.txt  control expression file (n x p tab-sep real mat)"
      << std::endl;
  std::cerr << " Y.txt  case expression file (n x p tab-sep real mat)"
      << std::endl;
  std::cerr << " K      number of components (default: " << args.K << ")"
      << std::endl;
  std::cerr << " out    output files' header (default: " << args.output << ")"
      << std::endl;
  std::cerr << " bmin   min value of beta (real)" << std::endl;
  std::cerr << " bmax   max value of beta (real)" << std::endl;
  std::cerr << " iter   max number of iterations (positive integer)"
      << std::endl;
  std::cerr << " tol    convergence condition (positive real)" << std::endl;
  std::cerr << " sample number of samples (positive integer)" << std::endl;
  std::cerr << " repeat number of repeated fits (positive integer)"
      << std::endl;
#ifdef ONLINE
  std::cerr << " minibatch size of mini-batch (default: " << args.minibatch
      << ")" << std::endl;
#endif
  std::cerr << std::endl;
  std::cerr << "" << std::endl;
  std::cerr << std::endl;
}

bool parse_args(const int argc, const char* argv[], prog_args_t& args_out)
{
  for (int j = 1; j < argc; ++j)
  {
    std::string curr = argv[j];

    if (curr == "-x" && ++j < argc)
    {
      args_out.x_file = argv[j];
    }
    else if (curr == "-y" && ++j < argc)
    {
      args_out.y_file = argv[j];
    }
    else if (curr == "-o" && ++j < argc)
    {
      args_out.output = argv[j];
    }
    else if (curr == "-i" && ++j < argc)
    {
      args_out.max_iter = get_param<size_t>(argv[j], "parsing iter");
    }
    else if (curr == "-s" && ++j < argc)
    {
      args_out.sample = get_param<size_t>(argv[j], "parsing #samples");
    }
    else if (curr == "-k" && ++j < argc)
    {
      args_out.K = get_param<size_t>(argv[j], "parsing K");
    }
    else if (curr == "-r" && ++j < argc)
    {
      args_out.repeat = get_param<size_t>(argv[j], "parsing repeat");
    }
    else if (curr == "-t" && ++j < argc)
    {
      args_out.tol = get_param<double>(argv[j], "parsing tol");
    }
    else if (curr == "-bmin" && ++j < argc)
    {
      args_out.bmin = get_param<double>(argv[j], "parsing beta min");
    }
    else if (curr == "-bmax" && ++j < argc)
    {
      args_out.bmax = get_param<double>(argv[j], "parsing beta max");
    }
    else if (curr == "-minibatch" && ++j < argc)
    {
      args_out.minibatch = get_param<size_t>(argv[j], "parsing minibatch");
    }

  }

  if (args_out.tol <= 0)
    return false;
  if (args_out.x_file.size() == 0 || args_out.y_file.size() == 0)
    return false;
  if (args_out.bmin >= args_out.bmax)
    return false;
  if (args_out.K < 1)
    return false;
  return true;
}

int main(const int argc, const char* argv[])
{
  prog_args_t args;
  if (!parse_args(argc, argv, args))
  {
    print_help(argv[0], args);
    exit(1);
  }

  expr_pair_vec_t data_vec;

  size_t _n, _p;
  if (!read_expr_pair_vec(args.x_file.c_str(), args.y_file.c_str(), data_vec,
      _n, _p))
  {
    err_msg("failed to read the data");
    return 1;
  }

  const int K = args.K;
  const int p = _p;

  double max_llik = 0.;
  boost::ptr_vector<func_expr_div_t> components;

  for (size_t r = 1; r <= args.repeat; ++r)
  {
    components.clear();
    for (size_t k = 0; k < K; ++k)
      components.push_back(new func_expr_div_t(p, args.bmin, args.bmax));

    double curr_llik;
#ifdef ONLINE
    curr_llik = fit_dpm_online<expr_pair_t, func_expr_div_t>(
        data_vec, components, args.max_iter, args.sample, args.tol,
        args.minibatch);
#else
    curr_llik = fit_dpm<expr_pair_t, func_expr_div_t>(data_vec, components,
        args.max_iter, args.sample, args.tol);
#endif

    if (r == 1 || curr_llik > max_llik)
    {
      max_llik = curr_llik;

      TLOG(" repeat = " << r << " max log-likelihood = " << max_llik);

      // output argmax membership
      std::ofstream ofs((args.output + ".argmax").c_str(), std::ios::out);
      for (size_t i = 0; i < data_vec.size(); ++i)
      {
        const expr_pair_t& expr = *data_vec.at(i).get();
        size_t argmax = 0;
        double maxval = components[argmax].score(expr);
        for (size_t k = 1; k < K; ++k)
        {
          double curr = components[k].score(expr);
          if (curr > maxval)
          {
            maxval = curr;
            argmax = k;
          }
        }
        ofs << i << "\t" << argmax << std::endl;
      }
      ofs.close();

      // output parameters
      std::ofstream ofs_mean((args.output + ".mean").c_str(), std::ios::out);
      std::ofstream ofs_sd((args.output + ".sd").c_str(), std::ios::out);

      for (size_t k = 0; k < K; ++k)
      {
        const func_expr_div_t::vec_t& mean = components[k].mean_vec();
        const func_expr_div_t::vec_t& var = components[k].var_vec();

        ofs_mean << k << "\t";
        vec_out(ofs_mean, mean);
        ofs_mean << std::endl;
        ofs_sd << k << "\t";
        vec_out(ofs_sd, var, sqrt_t());
        ofs_sd << std::endl;
      }

      ofs_mean.close();
      ofs_sd.close();
    }
  }

  return 0;
}

